Skip to content

[WIP] [feat] delta embedding dump upload to PAI Feature Store#603

Open
eric-gecheng wants to merge 14 commits into
alibaba:masterfrom
eric-gecheng:feat/delta_export_optimize
Open

[WIP] [feat] delta embedding dump upload to PAI Feature Store#603
eric-gecheng wants to merge 14 commits into
alibaba:masterfrom
eric-gecheng:feat/delta_export_optimize

Conversation

@eric-gecheng

Copy link
Copy Markdown
Collaborator

No description provided.

@eric-gecheng eric-gecheng added the codex-review Let Codex Review label Jul 19, 2026
@github-actions github-actions Bot removed the codex-review Let Codex Review label Jul 19, 2026
# Step zero is excluded from the delta publication contract.
logger.info("Skipping delta embedding dump at step %s.", global_step)
return None
if self._interval_steps is not None and global_step % self._interval_steps == 0:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After _sync_final_step() replaces each local step with the global MAX, this boundary may only have been dumped by the longest-running rank. For example, if ranks finish at steps 49 and 50, the shorter rank adopts 50 and skips here, so its trailing tracked rows are never written. Please reject uneven exhaustion for all multi-rank dump cadences, or base this skip on each rank’s actual last dump while preserving collective participation; add a 49/50 regression test.

should_dump = self._rank == 0 and time.monotonic() >= self._next_dump_time

any_failed = local_error is not None
if getattr(self, "_world_size", 1) > 1 and self._interval_secs is not None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This failure rendezvous is limited to minute cadence. In default step mode, rank zero can observe an async uploader failure immediately while peers skip the throttled marker check, raise alone, and leave peers hanging in the next training collective. Please synchronize the failure bit for every distributed FeatureStore cadence before raising; boundary dump failures need the same treatment.

dtype=torch.int32,
device=self._device,
)
torch.distributed.all_reduce(state, op=torch.distributed.ReduceOp.MAX)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minute cadence executes this CUDA all_reduce plus host-synchronizing .item() calls every training step, while the enabled pipeline also all-reduces batch availability. That adds at least two blocking collectives per batch for dumps minutes apart. Please coalesce timer/error state with the existing availability rendezvous or otherwise remove redundant hot-path synchronization.

table = pq.read_table(source)
source.seek(0)
self._validate_parquet_schema(table.schema, path)
for row in table.to_pylist():

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

read_table() + to_pylist() materializes every embedding element as Python objects, then copies rows into NumPy arrays and another sorted list before uploading. One large delta can therefore exhaust rank-zero training memory despite max_pending_steps. Please stream Parquet batches and use bounded/spilled deduplication so memory is capped independently of one dump’s size.

with self._condition:
if self._aborting:
return
self._add_discovered_steps_locked()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The worker calls _add_discovered_steps_locked() on every poll. Because committed dumps/manifests are retained, each poll rescans all historical steps, rereads success/manifest JSON and Parquet schemas, and does so while holding _condition, so submit()/close() latency grows with job history. Please reconcile once at startup and incrementally track later submissions/new entries.

single_rank_path = os.path.join(
output_dir, f"{file_prefix}_step_{global_step}.parquet"
)
if os.path.isfile(single_rank_path):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These canonical paths are sampled without validating their hashes/generation against the selected success marker. A restored run can replace the same step’s canonical files before the new generation commits, making this tool compare uncommitted local rows with the previous remote commit. Please verify the selected files against success["shards"], or fail explicitly when the committed bytes are no longer available.

@@ -0,0 +1 @@
feature_store_py @ https://gecheng-rec.oss-accelerate-overseas.aliyuncs.com/software/feature_store_py-2.2.7-py3-none-any.whl

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please source this credential-bearing SDK from an organization-controlled registry/release and pin its SHA-256 hash. A mutable wheel URL in a personal bucket makes dependency compromise equivalent to arbitrary code execution with the training job’s FeatureStore credentials.

"feature_store_config.region must not be empty "
"(it may come from ALIBABA_CLOUD_REGION)"
)
endpoint_url = endpoint if "://" in endpoint else f"//{endpoint}"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This only rejects URI userinfo; an arbitrary plaintext or attacker-controlled endpoint is still passed to the SDK together with cloud and FeatureDB credentials. Please require HTTPS and trusted FeatureStore endpoints (or an explicit trusted custom-endpoint opt-in), and reject unexpected path/query/fragment components.

Comment thread tzrec/protos/train.proto
// Existing FeatureStore entity associated with an automatically created
// DynamicEmbedding FeatureView.
required string feature_entity_name = 8;
// Explicit FeatureDB version for this incremental training run. It must be

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This says the version must be provisioned before upload, but the uploader contract is AUTO_CREATE_ON_FIRST_DELTA_MERGE and the implementation deliberately performs no version precheck. Please document that the explicit non-default version is created by the first MERGE.

@github-actions

github-actions Bot commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

Static review summary (tests/builds intentionally not run):

  • Distributed correctness blockers: step-cadence final flushing can drop trailing rows from a shorter rank, and rank-local dump/uploader failures can leave peers hanging.
  • Scaling risks: minute cadence adds redundant blocking collectives to every batch, whole steps are materialized in Python memory before upload, and the worker repeatedly rescans retained history.
  • Safety and verification: tighten SDK provenance and endpoint validation, and bind readback samples to the committed shard hashes/generation.
  • Documentation: the version-provisioning comment contradicts the implementation. Please also add user-facing docs and a runnable config covering installation, credential environment variables (including the security token), shared-POSIX outbox requirements, CUDA/sharding and equal-batch constraints, restart/single-writer behavior, and the readback command.

Detailed actionable findings are inline.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant